from transformers import pipeline
from typing import List, Dict
from src.base.agent import Agent
from zhipuai import ZhipuAI

MODEL = "glm-4-plus"
client = ZhipuAI(api_key='652056ee12c005f36b61ba8df67274b2.US6dlqS8ckjtuQv3')

class SummarizationAgent:
    def __init__(self):
        self.agent = Agent(
            name="Content Summarizer",
            instructions="""You are a research paper summarization expert. Your task is to:
            1. Extract key findings and methodologies
            2. Identify major contributions
            3. Maintain technical accuracy
            4. Create concise yet comprehensive summaries"""
        )

    def summarize(self, papers: List[Dict], min_length: int = 100) -> Dict:
        """生成论文集合的整体综述"""
        try:
            # 分批处理论文
            batch_size = 5
            batches = [papers[i:i+batch_size] for i in range(0, len(papers), batch_size)]
            all_summaries = []
            key_findings = []
            
            for batch in batches:
                # 合并批次中的论文内容
                batch_text = ""
                for paper in batch:
                    batch_text += f"Title: {paper['title']}\n"
                    batch_text += f"Content: {paper.get('content', '')}\n\n"

                # 生成批次综述
                prompt = f"""Please analyze the following research papers and generate:
                1. A concise summary for each paper
                2. Key findings and contributions
                3. Methodological highlights
                4. Technical innovations
                
                Format your response as:
                SUMMARIES:
                [One paragraph summary per paper]

                KEY_FINDINGS:
                - [Bullet points of key findings]

                Papers to analyze:
                {batch_text}"""

                response = client.chat.completions.create(
                    model=MODEL,
                    messages=[{"role": "system", "content": "You are a research paper synthesis expert."}, 
                              {"role": "user", "content": prompt}],
                    temperature=0.2,
                    max_tokens=2000
                )
                
                if not response.choices:
                    continue
                    
                content = response.choices[0].message.content
                
                # 解析响应内容
                parts = content.split("\nKEY_FINDINGS:")
                summaries = parts[0].replace("SUMMARIES:", "").strip()
                findings = parts[1].strip().split("\n") if len(parts) > 1 else []
                
                all_summaries.append(summaries)
                key_findings.extend([f.strip('- ') for f in findings if f.strip()])
            
            return {
                "overall_summary": "\n\n".join(all_summaries),
                "key_findings": key_findings
            }
            
        except Exception as e:
            print(f"生成综述失败: {str(e)}")
            return {"error": f"生成综述失败: {str(e)}"}

    def _extract_key_findings(self, text: str) -> List[str]:
        """使用智谱API提取关键发现"""
        try:
            prompt = f"""Please analyze the following research paper content and extract the key findings. Return them as a list of concise bullet points.

            Content to analyze:
            {text}"""

            response = client.chat.completions.create(
                model=MODEL,
                messages=[{"role": "system", "content": "You are a research paper analysis expert."}, 
                          {"role": "user", "content": prompt}],
                temperature=0.2
            )
            
            findings = response.choices[0].message.content.strip().split('\n')
            return [f.strip('- ') for f in findings if f.strip()]
        except Exception as e:
            print(f"提取关键发现失败: {str(e)}")
            return ["无法提取关键发现"]
